# -*- coding: utf-8 -*-
# =============================================================================
#     FileName:
#         Desc: 注意在MAC环境上，multiprocessing.Queue()存在问题
#       Author: GGA
#        Email:
#     HomePage:
#      Version: 1.0.1
#   LastChange: 2020-12-20
#      History:
# =============================================================================

import datetime
import multiprocessing as mp
import random
import time
import traceback
import uuid


class TimeoutException(Exception):
    def __init__(self, msg):
        super().__init__(msg)
        self.msg = msg

    def __str__(self):
        return self.msg

    def __repr__(self):
        return self.msg


class TimeoutTaskResult(object):
    def __init__(self, result=None, exception=None, traceback_exc=None):
        self.result = result
        self.exception = exception
        self.traceback_exc = traceback_exc


class TimeoutTask(object):
    def __init__(self, timeout, fn=None, args=None, kwargs=None):
        """
        带超时控制的任务类
        @param timeout: 超时时间,以秒为单位
        @param fn: 目标方法
        @param args: 目标参数
        @param kwargs: 目标参数
        """
        super().__init__()
        self.task_id = str(uuid.uuid4())
        if args is None:
            args = ()
        if kwargs is None:
            kwargs = {}
        self.fn = fn
        self.args = args
        self.kwargs = kwargs
        self.timeout = timeout
        self.timeout_time = time.time() + timeout
        self.result_queue = mp.Queue()
        self.daemon = True
        self.task_process = None

    @classmethod
    def _task_worker(cls, result_queue: mp.Queue, fn, args, kwargs):
        try:
            fn_result = fn(*args, **kwargs)
            task_result = TimeoutTaskResult(result=fn_result)
            result_queue.put(task_result)
        except Exception as ex:
            task_result = TimeoutTaskResult(exception=ex, traceback_exc=traceback.format_exc())
            result_queue.put(task_result)

    def start_task(self):
        self.task_process = mp.Process(
            target=self._task_worker,
            args=(self.result_queue, self.fn, self.args, self.kwargs)
        )
        self.task_process.daemon = True
        self.task_process.start()
        self.timeout_time = time.time() + self.timeout

    def check_task_running(self):
        """
        检查任务运行状态，如果任务正在执行且已经执行超时，则KILL任务并抛出超时异常到结果队列。
        """
        if self.task_process.is_alive():
            if self.timeout_time < time.time():
                msg = "The execution time exceeds {} seconds,timeout!".format(self.timeout)
                task_result = TimeoutTaskResult(
                    exception=TimeoutException(msg=msg),
                )
                self.result_queue.put(task_result)
                self.task_process.terminate()
                return False
            else:
                return True
        else:
            return False

    def get_task_result(self, wait_task_done: bool = False):
        """
        获取任务执行结果
        @param wait_task_done: 等待任务完成
        """
        while wait_task_done and self.check_task_running():
            time.sleep(0.001)
        if self.check_task_running():
            return TimeoutTaskResult(exception=Exception("Task is running!"))
        # 等待3毫秒，以保证读取到result_queue中数据
        time.sleep(0.003)
        if not self.result_queue.empty():
            return self.result_queue.get_nowait()
        else:
            return TimeoutTaskResult(exception=Exception("Unknown exception!"))


class TimeoutTaskPool(object):
    def __init__(self, max_workers):
        self.max_workers = max_workers
        self.task_items = dict()

    def add_task(self, timeout, fn=None, args=None, kwargs=None):
        task_item = TimeoutTask(
            timeout=timeout, fn=fn,
            args=args, kwargs=kwargs
        )
        self.task_items[task_item.task_id] = task_item
        return task_item.task_id

    def run_tasks(self):
        running_task_ids = set()
        waiting_task_ids = set(self.task_items.keys())
        tmp_task_ids = set()
        while True:
            tmp_task_ids.clear()
            for task_id in running_task_ids:
                task_item: TimeoutTask = self.task_items[task_id]
                if not task_item.check_task_running():
                    tmp_task_ids.add(task_id)
            for task_id in tmp_task_ids:
                running_task_ids.remove(task_id)
            tmp_task_ids.clear()
            for task_id in waiting_task_ids:
                if len(running_task_ids) > self.max_workers:
                    break
                task_item: TimeoutTask = self.task_items[task_id]
                task_item.start_task()
                running_task_ids.add(task_id)
                tmp_task_ids.add(task_id)
            for task_id in tmp_task_ids:
                waiting_task_ids.remove(task_id)
            if len(running_task_ids) == 0 and len(waiting_task_ids) == 0:
                break
            else:
                time.sleep(0.001)
            pass
        task_results = dict()
        for task_item in self.task_items.values():
            task_item: TimeoutTask = task_item
            task_results[task_item.task_id] = task_item.get_task_result(wait_task_done=True)
        return task_results


class TimeoutTaskDemo(object):
    @classmethod
    def print_info(cls, msg):
        print("[{}] {}".format(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S:%f"), msg))

    @classmethod
    def demo_task(cls, task_name, sleep_seconds):
        cls.print_info("start task {} with sleep {} seconds".format(task_name, sleep_seconds))
        time.sleep(sleep_seconds)
        cls.print_info("stop task {} with sleep {} seconds".format(task_name, sleep_seconds))
        return task_name, sleep_seconds

    @classmethod
    def show_task_result(cls, task_id: str, task_result: TimeoutTaskResult):
        task_msg = ""
        if task_result.result is not None:
            task_name, sleep_seconds = task_result.result
            task_msg += "use {} seconds.".format(
                sleep_seconds
            )
        else:
            task_msg += "none."
        if task_result.exception is not None:
            task_msg += " exception: {}".format(str(task_result.exception))
        cls.print_info("task_id: {}, task_result: {}".format(task_id, task_msg))

    @classmethod
    def demo(cls):
        cls.print_info("开始TimeoutTask演示")
        task_item = TimeoutTask(
            timeout=2,
            fn=cls.demo_task,
            args=("task1", 3)
        )
        task_item.start_task()
        task_result = task_item.get_task_result(wait_task_done=True)
        cls.show_task_result(task_id=task_item.task_id, task_result=task_result)
        cls.print_info("开始TimeoutTaskPool演示")
        ttp = TimeoutTaskPool(max_workers=2)
        for task_index in range(0, 10):
            ttp.add_task(
                timeout=3.3,
                fn=cls.demo_task,
                args=("task_{}".format(task_index), random.randint(1, 5))
            )
        task_results = ttp.run_tasks()
        for task_id in task_results.keys():
            task_result = task_results[task_id]
            cls.show_task_result(task_id=task_id, task_result=task_result)
        pass


if __name__ == '__main__':
    TimeoutTaskDemo.demo()
